#include "particle_renderer.h"
#include <GL/glew.h>
#include "engine/utils/file_utils.h"

namespace kanon
{
    ParticleBatch::~ParticleBatch()
    {
        glDeleteVertexArrays(1, &vao);
        glDeleteBuffers(1, &vbo);
        glDeleteBuffers(1, &ibo);
    }

    ParticleBatch::ParticleBatch(uint32_t max_particles)
    {
        max_vertices = max_particles * 4;
        max_indices = max_particles * 6;
        num_vertices = 0;
        num_indices = 0;
        
        vertex_base = nullptr;
        index_base = nullptr;

        // create vao and vbo
        glGenVertexArrays(1, &vao);
        glBindVertexArray(vao);

        glGenBuffers(1, &vbo);
        glBindBuffer(GL_ARRAY_BUFFER, vbo);
        glBufferData(GL_ARRAY_BUFFER, max_vertices * sizeof(ParticleVertex), nullptr, GL_DYNAMIC_DRAW);

        // set vertex attributes
        glVertexAttribPointer(0, 2, GL_FLOAT, GL_FALSE, sizeof(ParticleVertex), (void*)offsetof(ParticleVertex, pos));
        glEnableVertexAttribArray(0);
        glVertexAttribPointer(1, 4, GL_FLOAT, GL_FALSE, sizeof(ParticleVertex), (void*)offsetof(ParticleVertex, color));
        glEnableVertexAttribArray(1);

        glGenBuffers(1, &ibo);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo);
        glBufferData(GL_ELEMENT_ARRAY_BUFFER, max_indices * sizeof(uint32_t), nullptr, GL_DYNAMIC_DRAW);
    }

    void ParticleBatch::Begin()
    {
        // map buffer
        glBindVertexArray(vao);
        glBindBuffer(GL_ARRAY_BUFFER, vbo);
        vertex_base = (ParticleVertex*)glMapBuffer(GL_ARRAY_BUFFER, GL_WRITE_ONLY);
        glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ibo);
        index_base = (uint32_t*)glMapBuffer(GL_ELEMENT_ARRAY_BUFFER, GL_WRITE_ONLY);
        num_vertices = 0;
        num_indices = 0;
    }

    void ParticleBatch::End()
    {
        glUnmapBuffer(GL_ARRAY_BUFFER);
        vertex_base = nullptr;
        glUnmapBuffer(GL_ELEMENT_ARRAY_BUFFER);
        index_base = nullptr;
    }

    bool ParticleBatch::Submit(const Particle2D& p, const glm::vec4& color, const glm::vec2& base_position)
    {
        uint32_t particle_indices = 6;
        uint32_t particle_vertices = 4;
        if (num_indices + particle_indices > max_indices)
            return false;
        if (num_vertices + particle_vertices > max_vertices)
            return false;

        // copy vertex data
        ParticleVertex* data = vertex_base + num_vertices;
        glm::vec2 pos = p.position + base_position;
        data[0].pos = { pos.x - p.size.x / 2, pos.y - p.size.y / 2};
        data[0].color = color;

        data[1].pos = { pos.x - p.size.x / 2, pos.y + p.size.y / 2 };
        data[1].color = color;

        data[2].pos = { pos.x + p.size.x / 2, pos.y + p.size.y / 2 };
        data[2].color = color;

        data[3].pos = { pos.x + p.size.x / 2, pos.y - p.size.y / 2 };
        data[3].color = color;

        // copy index data
        uint32_t* index_data = index_base + num_indices;
        index_data[0] = num_vertices + 0;
        index_data[1] = num_vertices + 1;
        index_data[2] = num_vertices + 2;
        index_data[3] = num_vertices + 0;
        index_data[4] = num_vertices + 2;
        index_data[5] = num_vertices + 3;
 
        num_vertices += 4;
        num_indices += 6;
        return true;
    }

    bool ParticleBatch::Submit(const glm::vec2& pos, const glm::vec2& size, const glm::vec4& color)
    {
        uint32_t particle_indices = 6;
        uint32_t particle_vertices = 4;
        if (num_indices + particle_indices > max_indices)
            return false;
        if (num_vertices + particle_vertices > max_vertices)
            return false;

        // copy vertex data
        ParticleVertex* data = vertex_base + num_vertices;
        data[0].pos = { pos.x - size.x / 2, pos.y - size.y / 2};
        data[0].color = color;

        data[1].pos = { pos.x - size.x / 2, pos.y + size.y / 2 };
        data[1].color = color;

        data[2].pos = { pos.x + size.x / 2, pos.y + size.y / 2 };
        data[2].color = color;

        data[3].pos = { pos.x + size.x / 2, pos.y - size.y / 2};
        data[3].color = color;

        // copy index data
        uint32_t* index_data = index_base + num_indices;
        index_data[0] = num_vertices + 0;
        index_data[1] = num_vertices + 1;
        index_data[2] = num_vertices + 2;
        index_data[3] = num_vertices + 0;
        index_data[4] = num_vertices + 2;
        index_data[5] = num_vertices + 3;

        num_vertices += 4;
        num_indices += 6;
        return true;
    }

    ParticleRenderer::ParticleRenderer()
    {
        std::string vertex_shader_path = FileUtils::getShaderPath("particle_system.vs");
        std::string fragment_shader_path = FileUtils::getShaderPath("particle_system.fs");
        shader = Shader::FromShaderFiles(vertex_shader_path, fragment_shader_path);
        batch = std::make_shared<ParticleBatch>(50000);
    }

    void ParticleRenderer::SetCamera(const std::shared_ptr<Camera2d>& cam)
    {
        this->camera = cam;
    }

    void ParticleRenderer::Begin()
    {
        batch->Begin();
    }

    void ParticleRenderer::End()
    {
        this->Flush();
    }

    void ParticleRenderer::Flush()
    {
        batch->End();
        shader->Bind();

        if (this->camera)
        {
            glm::mat4 view = this->camera->GetViewMatrix();
            glm::mat4 proj = this->camera->GetProjectionMatrix();
            shader->set_matrix("vp", proj * view);
        }

        glDrawElements(GL_TRIANGLES, batch->num_indices, GL_UNSIGNED_INT, nullptr);
        shader->Unbind();
    }

    void ParticleRenderer::Render(const std::vector<Particle2D>& particles,
        const glm::vec4& color, const glm::vec2& base_position)
    {
        glEnable(GL_BLEND);
        glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
        for (auto& p : particles)
        {
            if (p.life < 0)
                continue;

            glm::vec4 c_color;
            c_color.r = color.r;
            c_color.g = color.g;
            c_color.b = color.b;
            c_color.a = p.life / p.total_life;
            // TODO: particle fade-out
            bool success = batch->Submit(p, c_color, base_position);
            if (!success)
            {
                // flush and restart
                this->Flush();
                this->Begin();
                batch->Submit(p, color, base_position);
            }
        }
    }
}
